package com.zoe.onelink.ratelimit.rule.register.impl;

import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.zoe.onelink.ratelimit.rule.register.RateLimitRuleRegister;
import io.opensergo.proto.fault_tolerance.v1.RateLimitStrategy;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * <p>标题: </p>
 * <p>描述: </p>
 * <p>版权: Copyright (c) 2022</p>
 * <p>公司: 智业软件股份有限公司</p>
 *
 * @version: 1.0
 * @author: ruansheng
 * @date: 2022-12-07
 */
public class SentinelRateLimitRuleRegister implements RateLimitRuleRegister {
    @Override
    public void registerRules(List<RateLimitStrategy> rateLimitStrategies) {
        if (rateLimitStrategies.isEmpty()) {
            return;
        }
        Map<RateLimitStrategy.MetricType, List<RateLimitStrategy>> ruleMap = rateLimitStrategies.stream()
                .collect(Collectors.groupingBy(RateLimitStrategy::getMetricType));
        // 暂时只考虑支持流控规则
        List<FlowRule> flowRules = ruleMap.get(RateLimitStrategy.MetricType.TYPE_REQUEST_AMOUNT).stream()
                .map(rateLimitStrategy -> {
                    double threshold = rateLimitStrategy.getThreshold() * 1.0 / rateLimitStrategy.getStatDuration();
                    FlowRule flowRule = new FlowRule();
                    // 资源名，资源名是限流规则的作用对象
                    flowRule.setResource(rateLimitStrategy.getName());
                    // 限流阈值类型，QPS 或线程数模式，这里使用 QPS 模式
                    flowRule.setGrade(RuleConstant.FLOW_GRADE_QPS);
                    // 限流阈值
                    flowRule.setCount(threshold < 1 ? 1 : threshold);
                    // 单机模式
                    flowRule.setClusterMode(false);
                    // 流控效果（直接拒绝 / 排队等待 / 慢启动模式），不支持按调用关系限流，默认直接拒绝
                    flowRule.setControlBehavior(RuleConstant.CONTROL_BEHAVIOR_DEFAULT);
                    return flowRule;
                }).collect(Collectors.toList());
        if (!flowRules.isEmpty()) {
            FlowRuleManager.loadRules(flowRules);
        }
    }

}
